#include "CryptoAgent.h"

#include <openssl/aes.h>
#include <openssl/des.h>
#include <openssl/rc4.h>

CryptoAgent::CryptoAgent() :
	_type(CipherType::None)
{
}

void CryptoAgent::setPassword(CipherType::Enum type, QString baseStr)
{
	_pass = baseStr;
	_type = type;

	int sz = _pass.size();
	if(sz == 0) return;

	// Normalize pass
	switch(_type)
	{
	case CipherType::AES: // 128 bit
		if(sz >= 16) return;
		while(_pass.size() < 16){
			_pass.append(_pass.at(_pass.size()%sz));
		}
		break;
	case CipherType::DES: // 64 bit
		if(sz >= 8) return;
		while(_pass.size() < 8){
			_pass.append(_pass.at(_pass.size()%sz));
		}
		break;
	//case CipherType::RC4: // 40 - 2048 bit
	//	if(sz >= 8) return;
	//	while(_pass.size() < 8){
	//		_pass.append(_pass.at(_pass.size()%sz));
	//	}
	//	break;
	}
}

QString CryptoAgent::encrypt(QString data)
{
	if(_pass.size() == 0) return data;
	switch(_type)
	{
	case CipherType::AES:
		return encryptAES(data, true);
	case CipherType::DES:
		return encryptDES(data, true);
	case CipherType::RC4:
		return encryptRC4(data, true);
	case CipherType::None:
		return data;
	}
}

QString CryptoAgent::decrypt(QString data)
{
	if(_pass.size() == 0) return data;
	switch(_type)
	{
	case CipherType::AES:
		return encryptAES(data, false);
	case CipherType::DES:
		return encryptDES(data, false);
	case CipherType::RC4:
		return encryptRC4(data, false);
	case CipherType::None:
		return data;
	}
}

QString CryptoAgent::encryptAES(QString data, bool encrypt)
{
	int sz = data.size();
	if(sz == 0) return "";

	unsigned char* outdata = new unsigned char[sz];
	unsigned char* ptr = (unsigned char*) data.toAscii().data();
	unsigned char ivec[] = "(unsigned char*)";

	AES_KEY key;
	AES_set_encrypt_key((unsigned char*) _pass.toAscii().data(), 128, &key);
	int num = 0;
	AES_cfb128_encrypt(ptr, outdata, sz, &key, ivec, &num, encrypt ? AES_ENCRYPT : AES_DECRYPT);

	QByteArray array;
	for(int i=0;i<sz;i++) array.append(outdata[i]);
	QString str(array);
	delete outdata;

	return str;
}

QString CryptoAgent::encryptDES(QString data, bool encrypt)
{
	int sz = data.size();
	if(sz == 0) return "";

	unsigned char* outdata = new unsigned char[sz];
	unsigned char* ptr = (unsigned char*) data.toAscii().data();
	DES_cblock Key2;
	DES_key_schedule schedule;

	/* Prepare the key for use with DES_cfb64_encrypt */
	memcpy( Key2, (unsigned char*) _pass.toAscii().data(), 8);
	DES_set_odd_parity( &Key2 );
	DES_set_key_checked( &Key2, &schedule );

	/* Encryption occurs here */
	int num = 0;
	DES_cfb64_encrypt( ptr, outdata, sz, &schedule, &Key2, &num, encrypt ? DES_ENCRYPT : DES_DECRYPT);

	QByteArray array;
	for(int i=0;i<sz;i++) array.append(outdata[i]);
	QString str(array);
	delete outdata;

	return str;
}

QString CryptoAgent::encryptRC4(QString data, bool encrypt)
{
	int sz = data.size();
	if(sz == 0) return "";
	int sz_or = sz;
	while(sz%16){data.append(" ");sz++;}

	unsigned char* outdata = new unsigned char[sz];
	unsigned char* ptr = (unsigned char*) data.toAscii().data();

	RC4_KEY key;
	int len = _pass.size();
	RC4_set_key(&key, len, (unsigned char*) _pass.toAscii().data());

	RC4(&key, sz, ptr, outdata);

	QByteArray array;
	for(int i=0;i<sz_or;i++) array.append(outdata[i]);
	QString str(array);
	delete outdata;

	return str;
}

void CryptoAgent::encryptNote(Note* note)
{
	note->setLogin( encrypt(note->getLogin()) );
	note->setPass( encrypt(note->getPass()) );
	note->setComment( encrypt(note->getComment()) );
	note->setUrl( encrypt(note->getUrl()) );
}

void CryptoAgent::decryptNote(Note* note)
{
	note->setLogin( decrypt(note->getLogin()) );
	note->setPass( decrypt(note->getPass()) );
	note->setComment( decrypt(note->getComment()) );
	note->setUrl( decrypt(note->getUrl()) );
}
